import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import hilbert
from scipy.interpolate import griddata

# 定义 x 的取值范围
x = np.linspace(-10, 10, 400)

# 选择不同的 c 值
c_values = np.arange(1,15,2)

# 存储所有的 (x, y) 坐标和对应的希尔伯特变换模值
all_x = []
all_y = []
all_hilbert_mod = []

# 循环遍历不同的 c 值
complex_x = np.asarray(x, dtype=complex)
c_values_clear = np.arange(1,15,0.1)
for c in c_values_clear:
    # 计算 y 的值
    try:
        y = np.sqrt(c**2 + 2*c*complex_x)
        y = y.real
    except:
        y=0
        pass
    y_combined = np.hstack((y[::-1], -y))
    x_combined = np.hstack((x[::-1], x))

    # 进行希尔伯特变换并取模
    hilbert_transform = hilbert(y_combined)
    hilbert_mod = np.abs(hilbert_transform) #this is mod, the y is theta.

    all_x.extend(x_combined)
    all_y.extend(y_combined)
    all_hilbert_mod.extend(hilbert_mod)

# 转换为 numpy 数组
all_x = np.array(all_x)
all_y = np.array(all_y)
all_hilbert_mod = np.array(all_hilbert_mod)

# 创建网格
xi = np.linspace(all_x.min(), all_x.max(), 500)
yi = np.linspace(all_y.min(), all_y.max(), 500)
XI, YI = np.meshgrid(xi, yi)

# 插值
ZI = griddata((all_x, all_y), all_hilbert_mod, (XI, YI), method='cubic')

# 绘制二维底板
plt.figure(figsize=(10, 6))
plt.imshow(ZI, cmap='viridis', aspect='auto',
           extent=[all_x.min(), all_x.max(), all_y.min(), all_y.max()], origin='lower', vmin=0,vmax=50)
plt.colorbar(label='Hilbert Transform Modulus')

# 再次循环绘制曲线簇
# 转换为复数数组
complex_x = np.asarray(x, dtype=complex)
for c in c_values:
    # 计算 y 的值
    try:
        y = np.sqrt(c**2 + 2*c*complex_x)
    except:
        y = 0
        pass
    y_combined = np.hstack((y[::-1], -y))
    x_combined = np.hstack((x[::-1], x))
    # 绘制曲线
    plt.plot(x_combined, y_combined, color='red', linewidth=1)

plt.title('2D Plot of Hilbert Transform Modulus with Curve Family')
plt.xlabel('x')
plt.ylabel('y')
plt.show()
